#include "RayMarchingProcess.h"
#include "Utils/LogUtil.h"
#include "Render/RenderState.h"
#include "Core/CoreGlobals.h"
#include "Game/Scene/Entity.h"
#include "Game/Component/Model.h"
#include "Render/Buffer/VertexBuffer.h"
#include "Render/Buffer/IndexBuffer.h"
#include "Render/Buffer/CustomBuffer.h"
#include "LocalShader.h"

void RayMarchingProcess::OnInit()
{
	mLightEntity = mGameScene->GetEntityByName("light");

	ENQUEUE_RENDER_COMMAND(BuildModel)([this](VkCommandBuffer commandBuffer)
		{
			
			mDescriptor = RHIDescriptor::Builder()
				.PushShader<RayMarchingVertexShader>()
				.PushShader<RayMarchingFragmentShader>()
				.Build();


			mColorRenderTarget = ColorRenderTarget::Builder()
				.SetUsage(VK_IMAGE_USAGE_COLOR_ATTACHMENT_BIT | VK_IMAGE_USAGE_SAMPLED_BIT)
				.SetStoreOp(VK_ATTACHMENT_STORE_OP_STORE)
				.SetFinalLayout(VK_IMAGE_LAYOUT_SHADER_READ_ONLY_OPTIMAL)
				.Build(this, "ColorRenderTarget");

			mDepthRenderTarget = DepthRenderTarget::Builder()
				.Build();

			mPipeline = GraphicsPipeline::Builder()
				.SetDescriptor(mDescriptor)
				.AddColorRenderTarget(mColorRenderTarget)
				.SetDepthRenderTarget(mDepthRenderTarget)
				.SetDstStageMask(VK_PIPELINE_STAGE_COLOR_ATTACHMENT_OUTPUT_BIT)
				.SetDstAccessMask(VK_ACCESS_COLOR_ATTACHMENT_WRITE_BIT)
				.SetSrcStageMask(VK_PIPELINE_STAGE_COLOR_ATTACHMENT_OUTPUT_BIT)
				.SetSrcAccessMask(VK_ACCESS_COLOR_ATTACHMENT_WRITE_BIT)
				.Build();

			mRenderPass = RHIRenderPass::Builder()
				.AddRenderTarget(mColorRenderTarget)
				.AddRenderTarget(mDepthRenderTarget)
				.AddPipeline(mPipeline)
				.SetSrcStageMask(VK_PIPELINE_STAGE_FRAGMENT_SHADER_BIT)
				.SetSrcAccessMask(VK_ACCESS_SHADER_READ_BIT)
				.SetDstStageMask(VK_PIPELINE_STAGE_FRAGMENT_SHADER_BIT)
				.SetDstAccessMask(VK_ACCESS_SHADER_READ_BIT)
				.Build(this, "RenderPass");

			mFrameBuffer = RHIFrameBuffer::Builder()
				.SetRenderPass(mRenderPass)
				.Build();
		});
}

void RayMarchingProcess::OnTick(float deltaTime)
{
	PushConsts pushConsts;

	glm::mat4 currView = mGameScene->GetCamera().GetView();
	glm::mat4 currProj = mGameScene->GetCamera().GetProj();

	pushConsts.viewProj = currView * currProj;
	pushConsts.viewInverse = glm::inverse(currView);
	pushConsts.projInverse = glm::inverse(currProj);
	pushConsts.lightPos = mLightEntity->GetTransform().GetLocalPosition();

	ENQUEUE_RENDER_COMMAND(RecordCommand)([this, pushConsts](VkCommandBuffer commandBuffer)
		{
			VkRenderPassBeginInfo renderPassBeginInfo = mFrameBuffer->GetRenderPassBeginInfo();

			vkCmdBeginRenderPass(commandBuffer, &renderPassBeginInfo, VK_SUBPASS_CONTENTS_INLINE);

			vkCmdSetViewport(commandBuffer, 0, 1, &GRenderState->GetViewport());
			vkCmdSetScissor(commandBuffer, 0, 1, &GRenderState->GetScissor());

			vkCmdBindPipeline(commandBuffer, VK_PIPELINE_BIND_POINT_GRAPHICS, mPipeline->GetPipeline());

			PushConsts factData = pushConsts;
			vkCmdPushConstants(commandBuffer, mPipeline->GetPipelineLayout(), VK_SHADER_STAGE_FRAGMENT_BIT, 0, sizeof(PushConsts), &factData);

			vkCmdDraw(commandBuffer, 3, 1, 0, 0);

			vkCmdEndRenderPass(commandBuffer);
		});
}

void RayMarchingProcess::OnResize()
{
	mColorRenderTarget = ColorRenderTarget::Builder()
		.SetUsage(VK_IMAGE_USAGE_COLOR_ATTACHMENT_BIT | VK_IMAGE_USAGE_SAMPLED_BIT)
		.SetStoreOp(VK_ATTACHMENT_STORE_OP_STORE)
		.SetFinalLayout(VK_IMAGE_LAYOUT_SHADER_READ_ONLY_OPTIMAL)
		.Build(this, "ColorRenderTarget");

	mDepthRenderTarget = DepthRenderTarget::Builder()
		.Build();

	mFrameBuffer = RHIFrameBuffer::Builder()
		.AddRenderTarget(mColorRenderTarget)
		.AddRenderTarget(mDepthRenderTarget)
		.SetRenderPass(mRenderPass)
		.Build();
}
